package smoothing;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import parser.Parser;

/**
 * Estimates probabilities according to training using Held-out smoothing
 */
public class HeldOut implements Smoothing {

	private long numEventsHeldOut; // number of events in held-out set
	private Map<String, Integer> trainObsVoc; // maps each word in train set to its number of instances
	private Map<String, Integer> wordCtrHeldOut; // maps each word in train set to its number of instances in held-out set
	private Map<Integer, List<String>> wordsByInst; // maps each instance number to all of the words that appear as many times in train set
 
	public HeldOut(String trainFileName, String heldoutFileName) throws IOException {
		
		// count number of events in held-out set
		numEventsHeldOut = Parser.numEvents(heldoutFileName);

		// count number of instances for each word in train set
		trainObsVoc = Parser.genObsVoc(trainFileName);
		
		// count instances in held-out set for each word in train set
		wordCtrHeldOut = new HashMap<String, Integer>();
		countTrainInHeldout(trainFileName, heldoutFileName);
		
		// arrange words in train set by number of instances
		wordsByInst = Parser.wordsByInstNum(trainFileName);
	}
	
	private void countTrainInHeldout(String trainFileName, String heldoutFileName) throws IOException {

		// get all different words in train set
		Map<String, Integer> trainObsVoc = Parser.genObsVoc(trainFileName);

		// count instances for each word in held-out set
		Map<String, Integer> heldoutObsVoc = Parser.genObsVoc(heldoutFileName);
		
		// for each word in train set, count its instances in held-out set
		for (String currWord : trainObsVoc.keySet()) {
			
			// if word appears in held-out set, get its number of instances in held-out set
			if (heldoutObsVoc.containsKey(currWord)) {
				wordCtrHeldOut.put(currWord, heldoutObsVoc.get(currWord));
				continue;
			}
			
			// word does not appear in held-out set
			wordCtrHeldOut.put(currWord, 0);
		}
	}
	
	/**
	 * Calculates the MLE for given word, using Held-out smoothing
	 */
	@Override
	public double mleModel(String word) {
		
		// get number of instances of given word in train set
		int numInstTrain = trainObsVoc.containsKey(word) ? trainObsVoc.get(word) : 0;
		
		// in case of a a word that is not in the training set
		if (numInstTrain == 0)
			return mleModel();
		
		// get all words from train set with this number of instances in held-out set
		List<String> words = wordsByInst.get(numInstTrain);
		
		// sum number of instances in held-out for these words
		int sumInsances = 0;
		
		for (String currWord : words) {
			sumInsances += wordCtrHeldOut.get(currWord);
		}
		
		// return estimated probability
		return ((double) sumInsances) / ((double) words.size() * numEventsHeldOut);
	}

	/**
	 * Calculates the MLE for a word not in the training set, using
	 * Held-out smoothing
	 * @return MLE for an unseen word according to Held-out smoothing
	 */
	public double mleModel() {
		
		// sum number of instances in held-out that was not observed in the train
		int sumInsancesInHOButNotInTrain = 0;
		for (Map.Entry<String, Integer> currEntry : wordCtrHeldOut.entrySet()) {
			
			if ( ! trainObsVoc.containsKey(currEntry.getKey()))
				sumInsancesInHOButNotInTrain += currEntry.getValue();
			}		
		
		// calculate the number of events that was not observed in the train
		long unseenEventsNum = Parser.VOCAB_SIZE - trainObsVoc.size();
		
		// return estimated probability
		return ((double) sumInsancesInHOButNotInTrain) / ((double) unseenEventsNum * numEventsHeldOut);
	}
	
	/**
	 * Returns true if the probabilities' sum is 1, otherwise false
	 * @return correctness of the smoothing
	 */
	@Override
	public boolean debug() {
		
		// number of unseen events
		long unseenEventsNum = Parser.VOCAB_SIZE - trainObsVoc.size();
		
		// sum the probabilities of the unseen events
		double unseenProbabilitiesSum = mleModel() * unseenEventsNum;
		
		// sum the probabilities of the observed events
		double obsEventsProbabilitiesSum = 0;
		
		for (Map.Entry<String, Integer> currEntry : trainObsVoc.entrySet()) {
		
			obsEventsProbabilitiesSum += mleModel(currEntry.getKey());
			}
		
		// sum of the probabilities in the method 
		double probabilitiesSum = unseenProbabilitiesSum + obsEventsProbabilitiesSum;
		
		// in case the total sum is 1 --> return true
		if (probabilitiesSum == 1)
			return true;
		
		// otherwise --> false
		return false;
	}
	
	/**
	 * Return a map with numOfWords and sumInsances of each frequency in the test set
	 */
	public Map<Integer, Map.Entry<Integer, Integer>> getInstancesDetails() {
		
		Map resultsMap = new HashMap<Integer, Map.Entry<Integer, Integer>>();
		
		// get the results for frequency 0:
		
		// calculate the number of events that was not observed in the train
		long unseenEventsNum = Parser.VOCAB_SIZE - trainObsVoc.size(); ///Nr
		
		// sum number of instances in held-out that was not observed in the train
		int sumInsances = 0;
		for (Map.Entry<String, Integer> currEntry : wordCtrHeldOut.entrySet()) {
			
			if ( ! trainObsVoc.containsKey(currEntry.getKey()))
				sumInsances += currEntry.getValue();
			}		
		
		// add the results to the map
		Map.Entry<Integer, Integer> unseenWordsEntry = new HashMap.SimpleEntry<Integer, Integer> ((int) unseenEventsNum, sumInsances);
		resultsMap.put(0, unseenWordsEntry);
		
		
		// fill the map with the rest of the results
		for (int frequency = 1; frequency < 10; frequency++)
		{
			// get all words from train set with this number of instances in held-out set
			List<String> words = wordsByInst.get(frequency);
			
			// sum number of instances in held-out for these words
			sumInsances = 0;
			for (String currWord : words) {
				sumInsances += wordCtrHeldOut.get(currWord);
			}
			
			// create and add a new entry to the map
			Map.Entry<Integer, Integer> newEntry = new HashMap.SimpleEntry<Integer, Integer> (words.size(), sumInsances);
			resultsMap.put(frequency, newEntry);		
		}
		
		return resultsMap;
	}
	
	/**
	 * Return the expected frequency for given frequency, using Held-out smoothing
	 */
	@Override
	public double expectedFrequency(int frequency) {
		
		double estimatedProbability;
		
		// in case of frequency 0
		if (frequency == 0)
		{
			estimatedProbability = mleModel();
		}
		else
		{
			// get all words from train set with this frequency in held-out set
			List<String> words = wordsByInst.get(frequency);
			
			// sum number of instances in held-out for these words
			int sumInsances = 0;
			for (String currWord : words) {
				sumInsances += wordCtrHeldOut.get(currWord);
			}
			
			// the estimated probability
			estimatedProbability = ((double) sumInsances) / ((double) words.size() * numEventsHeldOut);
		}
		
		// calculate the size of the training set
		int trainingSize = 0;
		
		for (Map.Entry<String, Integer> currEntry : trainObsVoc.entrySet()) {
			
			trainingSize += currEntry.getValue();
			}
		
		// return the expected frequency
		return estimatedProbability * trainingSize;
	}


}


